from collections import OrderedDict
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any, Iterable, Iterator, Mapping, TypeVar

from zkl_pyutils_registries import SpecialKey, UNSET

from .tag import Tag

context_stack_var = ContextVar[list[tuple[Tag, Any]]]("convertor_context_stack", default=None)


def get_context_stack():
    context_stack = context_stack_var.get()
    if context_stack is None:
        context_stack = []
        context_stack_var.set(context_stack)
    return context_stack


@contextmanager
def context(entries: Mapping[Tag, Any] | Iterable[tuple[Tag, Any]]):
    context_stack = get_context_stack()
    entries = tuple(entries.items()) if isinstance(entries, Mapping) else tuple(entries)
    try:
        # noinspection PyTypeChecker
        context_stack.extend(entries)
        yield
    finally:
        for entry in reversed(entries):
            assert context_stack.pop() == entry


def iter_context() -> Iterator[tuple[Tag, Any]]:
    context_stack = get_context_stack()
    if context_stack is not None:
        yield from reversed(context_stack)


AnyKey = TypeVar("AnyKey")
AnyQuery = TypeVar("AnyQuery")
AnyMatch = TypeVar("AnyMatch")


def iter_context_tag_value(tag: Tag[AnyQuery, AnyKey, AnyMatch]) -> Iterator[AnyKey]:
    for context_tag, context_key in iter_context():
        if context_tag == tag:
            yield context_key


def get_context_tag_value(
    tag: Tag[AnyQuery, AnyKey, AnyMatch],
    default: AnyKey | SpecialKey = UNSET,
) -> AnyKey | SpecialKey:
    for value in iter_context_tag_value(tag):
        return value
    return default


def get_context_tags() -> OrderedDict[Tag, Any]:
    keys = set()
    tags = []
    for key, value in iter_context():
        if key in keys:
            continue
        keys.add(key)
        tags.append((key, value))
    return OrderedDict(tags)
